package com.boge.ai.milvus;

import io.milvus.client.MilvusClient;
import io.milvus.client.MilvusServiceClient;
import io.milvus.grpc.DataType;
import io.milvus.grpc.SearchResults;
import io.milvus.param.ConnectParam;
import io.milvus.param.IndexType;
import io.milvus.param.MetricType;
import io.milvus.param.collection.CreateCollectionParam;
import io.milvus.param.collection.FieldType;
import io.milvus.param.collection.HasCollectionParam;
import io.milvus.param.dml.InsertParam;
import io.milvus.param.dml.SearchParam;
import io.milvus.param.index.CreateIndexParam;
import io.milvus.response.SearchResultsWrapper;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

/**
 * Milvus 案例
 */
public class MilvusExample {

    public static void main(String[] args) {
        // 1. 测试对应的数据集是否存在，测试Milvus的服务是否能够正常的连接
        // milvusConnectTest();
        // 2.创建对应的数据集
        // createCollection();
        // 3.创建向量索引
        // createVectorIndex();
        // 4. 插入数据
        // insertVectorData();
        // 5. 查询向量数据
        searchVectorData();
    }

    /**
     * 测试服务是否正常
     */
    public static void milvusConnectTest(){
        // 服务的连接信息
        ConnectParam connectParam = ConnectParam
                .newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();

        MilvusClient milvusClient = new MilvusServiceClient(connectParam);
        System.out.println(milvusClient);
        // 数据集的配置信息
        HasCollectionParam collectionParam = HasCollectionParam
                .newBuilder()
                .withCollectionName("llm_rag")
                .build();
        System.out.println(milvusClient.hasCollection(collectionParam));
    }

    /**
     * 创建数据集【类似于MYSQL数据库中的表结构】
     */
    public static void createCollection(){
        // 服务的连接信息
        ConnectParam connectParam = ConnectParam
                .newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();

        MilvusClient milvusClient = new MilvusServiceClient(connectParam);
        // 对应的数据集的字段集合
        List<FieldType> list = new ArrayList<>();
        list.add(FieldType
                .newBuilder()
                .withName("id")
                .withDescription("主键")
                .withDataType(DataType.Int64)
                .withPrimaryKey(true)
                .withAutoID(false)
                .build());
        list.add(FieldType
                .newBuilder()
                .withName("vector")
                .withDescription("向量")
                .withDataType(DataType.FloatVector)
                .withDimension(128)
                .build());
        // 数据集的相关信息
        CreateCollectionParam createCollectionParam = CreateCollectionParam
                .newBuilder()
                .withCollectionName("milvus_collection")
                .withDescription("向量数据库的测试数据集")
                .withShardsNum(2)
                .withFieldTypes(list) // 设置对应的 字段信息
                .build();

        // 创建对应的 数据集
        System.out.println(milvusClient.createCollection(createCollectionParam));

    }

    /**
     * 创建向量索引
     */
    public static void createVectorIndex(){
        // 服务的连接信息
        ConnectParam connectParam = ConnectParam
                .newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();

        MilvusClient milvusClient = new MilvusServiceClient(connectParam);
        // 创建索引的参数
        CreateIndexParam createIndexParam = CreateIndexParam.newBuilder()
                .withCollectionName("milvus_collection")
                .withIndexName("vector")
                .withFieldName("vector")
                .withIndexType(IndexType.IVF_FLAT)
                .withMetricType(MetricType.L2) // 制定度量类型
                .withExtraParam("{\"nlist\":128}") // 根据索引类型提供额外参数，比如nlist
                .build();
        System.out.println(milvusClient.createIndex(createIndexParam));
    }

    /**
     * 向 向量数据库 添加数据
     */
    public static void insertVectorData(){
        // 服务的连接信息
        ConnectParam connectParam = ConnectParam
                .newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();

        MilvusClient milvusClient = new MilvusServiceClient(connectParam);
        // id的数据
        List<Long> idList = Arrays.asList(1l,2l);
        // vector的数据
        List<List<Float>> vectorList = Arrays.asList(
                Collections.nCopies(128,0.1f),
                Collections.nCopies(128,0.2f)
        );

        // 添加需要添加的数据的集合
        List<InsertParam.Field> list = new ArrayList<>();
        list.add(new InsertParam.Field("id",idList));
        list.add(new InsertParam.Field("vector",vectorList));
        InsertParam insertParam = InsertParam.newBuilder()
                .withCollectionName("milvus_collection")
                .withFields(list)
                .build();
        milvusClient.insert(insertParam);
    }

    /**
     * 向量数据的查询
     */
    public static void searchVectorData(){
        // 服务的连接信息
        ConnectParam connectParam = ConnectParam
                .newBuilder()
                .withHost("localhost")
                .withPort(19530)
                .build();

        MilvusClient milvusClient = new MilvusServiceClient(connectParam);
        // 定义的向量数据集
        List<List<Float>> searchVectorList = Collections.singletonList(Collections.nCopies(128,0.2f));
        // 查询的参数对象
        SearchParam searchParam = SearchParam.newBuilder()
                .withCollectionName("milvus_collection")
                .withMetricType(MetricType.L2)
                .withTopK(3) // 返回最接近的前三个记录
                .withVectorFieldName("vector")
                .withVectors(searchVectorList)
                .withOutFields(Arrays.asList("id","vector")) // 定义需要返回的字段
                .build();

        SearchResults data = milvusClient.search(searchParam).getData();
        if(data != null){
            SearchResultsWrapper resultsWrapper = new SearchResultsWrapper(data.getResults());
            resultsWrapper.getRowRecords().forEach(results ->{
                System.out.println(results);
            });
        }
    }
}
